package org.nicta.learn.factorization.PMF;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Random;

import org.nicta.colt.utils.ColtUtils;
import org.nicta.datasetLoader.MovieLens.MovieLensDataLoader;
import org.nicta.learn.optimization.LBFGS;

import cern.colt.matrix.DoubleFactory2D;
import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.jet.math.Functions;

public class PMF {
	private int dimension;

	DoubleMatrix2D U;
	DoubleMatrix2D V;
	Algebra algebra;
	
	public HashMap<Integer[], Double> ratings;

	public double lambdaUser = 10;
	public double lambdaItem = 10;
	public int userCount;
	public int itemCount;
	
	public PMF(int dimension, HashMap<Integer[], Double> ratings , int userCount, int itemCount) throws Exception{
		Random random = new Random();
		this.dimension = dimension;
		this.ratings = ratings;
		this.userCount = userCount;
		this.itemCount = itemCount;
		
		U = DoubleFactory2D.dense.random(dimension, userCount);
		V = DoubleFactory2D.dense.random(dimension,  itemCount);
		algebra = new Algebra();
	}
	
	
	public DoubleMatrix2D getUserGradient(DoubleMatrix2D U, DoubleMatrix2D V){
		Integer userId;
		Integer itemId;
		double value;
		double predicted;
		double error;
		DoubleMatrix2D userGradient = DoubleFactory2D.dense.make(dimension, userCount);
		for(Integer[] userItemPair: ratings.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = ratings.get(userItemPair);
			predicted = algebra.mult(U.viewColumn(userId-1), V.viewColumn(itemId-1));
			error = value - predicted;
			DoubleMatrix1D tempUserGradientColumn = userGradient.viewColumn(userId-1);
			DoubleMatrix1D tempItemColumn = V.viewColumn(itemId-1);
			DoubleMatrix1D temp = ColtUtils.multScalar(tempItemColumn, -error);
			tempUserGradientColumn.assign(temp,Functions.plus);	
		}
		userGradient.assign(ColtUtils.multScalar(U, lambdaUser),Functions.plus);
		return userGradient;
	}
	
	public DoubleMatrix2D getMovieGradient(DoubleMatrix2D U, DoubleMatrix2D V){
		Integer userId;
		Integer itemId;
		double value;
		double predicted;
		double error;
		DoubleMatrix2D movieGradient = DoubleFactory2D.dense.make(dimension, itemCount);
		for(Integer[] userItemPair: ratings.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = ratings.get(userItemPair);
			predicted = algebra.mult(U.viewColumn(userId-1), V.viewColumn(itemId-1));
			error = value - predicted;
			DoubleMatrix1D tempItemGradientColumn = movieGradient.viewColumn(itemId-1);
			DoubleMatrix1D tempUserColumn = U.viewColumn(userId-1);
			DoubleMatrix1D temp = ColtUtils.multScalar(tempUserColumn, -error);
			tempItemGradientColumn.assign(temp,Functions.plus);
		}
		movieGradient.assign(ColtUtils.multScalar(V, lambdaItem),Functions.plus);
		return movieGradient;
	}
	
	public double[] calculateTestError(HashMap<Integer[], Double> testRatings){
		double rmseCost = 0.0;
		double maeCost = 0.0;
		double predicted;
		double value;
		int userId;
		int itemId;
		double[] result = new double [2];
		double threshold = 0.5;
		double tp = 0;
		double total = 0;
		for(Integer[] userItemPair:testRatings.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = testRatings.get(userItemPair);
			predicted = algebra.mult(U.viewColumn(userId-1), V.viewColumn(itemId-1));
			//cutoff
//			if (predicted > 1) predicted = 1;
//			if (predicted < 1) predicted = 1;
//			System.out.format("Actual : %f Predicted: %f\n",value,predicted);
//			rmseCost += Math.pow((value - predicted),2);
//			maeCost += Math.abs(value - predicted);
			predicted = predicted < threshold ? 0 : 1;
			if(predicted == value){
				tp += 1;
			}
			total  += 1;
			
		}
//		result[0] = Math.sqrt(rmseCost*(1.0/testRatings.size()));
//		result[1] = maeCost*(1.0/testRatings.size());
		result[0] = tp / total;
		result[1] = tp / total;
		return result; 
	}
	
	public double calculateCost(DoubleMatrix2D U, DoubleMatrix2D V){
		double cost = 0;
		int userId;
		int itemId;
		double value;
		double predicted;
		double norm;
		for(Integer[] userItemPair:ratings.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = ratings.get(userItemPair);
			predicted = algebra.mult(U.viewColumn(userId-1), V.viewColumn(itemId-1));
			cost += Math.pow((value-predicted), 2);			
		}
		norm = (Math.pow(algebra.normF(U),2) * lambdaUser+ 
				Math.pow(algebra.normF(V),2)* lambdaItem)/2; 
		return cost/2+norm;
	}
	
	public void factorize() throws Exception{

		DoubleMatrix2D userGradient ;
		DoubleMatrix2D movieGradient;
	 
		double prev_error;
		
		double error = Double.MAX_VALUE;
		double epoch = 1;
		
		int flatUsersLength = userCount * dimension;
		int flatItemLength = itemCount * dimension;
		
		double[] diag = new double[flatUsersLength + flatItemLength];
		
		int[] iprint = {0,0};
		int[] iflag = {0};
		
		double[] variables = new double[flatUsersLength + flatItemLength];
		
		double[] derivatives = new double[flatUsersLength + flatItemLength];
		
		do{
			prev_error = error;
			userGradient = getUserGradient(U, V);
			movieGradient = getMovieGradient(U,V);
			
			error = calculateCost(U,V);
			System.out.println("Iteration : "+ epoch);			
			System.out.println("New Error : " + error);
			System.out.println();
//			double[] cost = calculateTestError();
//			System.out.format("RMSE : %f\t MSE: %f\n\n", cost[0],cost[1]);
			
			ColtUtils.mergeFlattenTo1D(U, V, variables);
			ColtUtils.mergeFlattenTo1D(userGradient, movieGradient, derivatives);
			
			LBFGS.lbfgs(variables.length, dimension, variables, error, derivatives,
					false, diag, iprint, 0.0001,
					1e-15, iflag);
			
			ColtUtils.assignByUnflattening(variables,U,V);
			epoch++;

		}while(Math.abs(error-prev_error) > 0.0001 && iflag[0] != 0);
	}
	
}
